#include <algorithm>
#include <cassert>
#include "BVH.hpp"

BVHAccel::BVHAccel(std::vector<Object*> p, int maxPrimsInNode,
                   SplitMethod splitMethod)
    : maxPrimsInNode(std::min(255, maxPrimsInNode)), splitMethod(splitMethod),
      primitives(std::move(p))
{
    time_t start, stop;
    time(&start);
    if (primitives.empty())
        return;

    root = recursiveBuild(primitives);

    time(&stop);
    double diff = difftime(stop, start);
    int hrs = (int)diff / 3600;
    int mins = ((int)diff / 60) - (hrs * 60);
    int secs = (int)diff - (hrs * 3600) - (mins * 60);

    printf(
        "\rBVH Generation complete: \nTime Taken: %i hrs, %i mins, %i secs\n\n",
        hrs, mins, secs);
}

struct SAHBucket {
    Bounds3 bounds;
    int count{};
};

BVHBuildNode* BVHAccel::recursiveBuild(std::vector<Object*> objects)
{
    BVHBuildNode* node = new BVHBuildNode();

    // Compute bounds of all primitives in BVH node
    Bounds3 bounds;
    for (int i = 0; i < objects.size(); ++i)
        bounds = Union(bounds, objects[i]->getBounds());
    if (objects.size() == 1) {
        // Create leaf _BVHBuildNode_
        node->bounds = objects[0]->getBounds();
        node->object = objects[0];
        node->left = nullptr;
        node->right = nullptr;
        return node;
    }
    else if (objects.size() == 2) {
        node->left = recursiveBuild(std::vector{objects[0]});
        node->right = recursiveBuild(std::vector{objects[1]});

        node->bounds = Union(node->left->bounds, node->right->bounds);
        return node;
    }
    else {
        Bounds3 centroidBounds;
        for (int i = 0; i < objects.size(); ++i)
            centroidBounds =
                Union(centroidBounds, objects[i]->getBounds().Centroid());
        int dim = centroidBounds.maxExtent();
        switch (dim) {
            case 0:
                std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                    return f1->getBounds().Centroid().x <
                           f2->getBounds().Centroid().x;
                });
                break;
            case 1:
                std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                    return f1->getBounds().Centroid().y <
                           f2->getBounds().Centroid().y;
                });
                break;
            case 2:
                std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                    return f1->getBounds().Centroid().z <
                           f2->getBounds().Centroid().z;
                });
                break;
        }

        int nBuckets = 12;
        std::vector<SAHBucket> buckets(nBuckets);

        for (int i = 0; i < objects.size(); ++i) {
            auto && offset = centroidBounds.Offset(objects[i]->getBounds().Centroid());
            // dim 可能取值是0,1,2 正好可以获取Vector3f的分量
            int b = nBuckets * offset[dim];
            if (b >= nBuckets){
                b = nBuckets - 1;
            }
            assert(b >= 0 && b <nBuckets);
            buckets[b].count++;
            buckets[b].bounds = Union(buckets[b].bounds, objects[i]->getBounds());
        }
        std::vector<double> cost(nBuckets - 1);
        for (int i = 0; i < nBuckets - 1; ++i) {
            Bounds3 b0, b1;
            int count0 = 0, count1 = 0;
            //第一个桶到[i]桶，所有的图元数与边界盒
            for (int j = 0; j <= i; ++j) {
                b0 = Union(b0, buckets[j].bounds);
                count0 += buckets[j].count;
            }
            //[i+1]桶到最后一个桶，所有的图元数与边界盒
            for (int j = i + 1; j < nBuckets; ++j) {
                b1 = Union(b1, buckets[j].bounds);
                count1 += buckets[j].count;
            }
            //bounds变量为所有图元的边界盒
            //通过面积模型计算，相交开销设为1
            cost[i] = 1 +
                      (count0 * b0.SurfaceArea() +
                       count1 * b1.SurfaceArea()) /
                      bounds.SurfaceArea();
        }

        //计算出最小消耗的切割位置与消耗的量
        double minCost = cost[0];
        int minCostSplitBucket = 0;
        for (int i = 1; i < nBuckets - 1; ++i) {
            if (cost[i] < minCost) {
                minCost = cost[i];
                minCostSplitBucket = i;
            }
        }
        int pmid = 0;
        for (int i = 0; i <= minCostSplitBucket; ++i) {
            pmid += buckets[i].count;
        }

        auto beginning = objects.begin();
        auto middling = objects.begin() + pmid;
        auto ending = objects.end();

        auto leftshapes = std::vector<Object*>(beginning, middling);
        auto rightshapes = std::vector<Object*>(middling, ending);

        assert(objects.size() == (leftshapes.size() + rightshapes.size()));

        node->left = recursiveBuild(leftshapes);
        node->right = recursiveBuild(rightshapes);

        node->bounds = Union(node->left->bounds, node->right->bounds);
    }

    return node;
}

Intersection BVHAccel::Intersect(const Ray& ray) const
{
    Intersection isect;
    if (!root)
        return isect;
    isect = BVHAccel::getIntersection(root, ray);
    return isect;
}

Intersection BVHAccel::getIntersection(BVHBuildNode* node, const Ray& ray) const
{
    std::array<int, 3> dirIsNeg = {
            (int)(ray.direction.x > 0),
            (int)(ray.direction.y > 0),
            (int)(ray.direction.z > 0)
    };
    auto && invDir = ray.direction_inv;
    if (!node->bounds.IntersectP(ray, invDir, dirIsNeg)){
        return Intersection();
    }
    // recursiveBuild建立的BVH满足以下设定：叶子节点的left和right同时为nullptr
    // 为了简化和健壮,左右节点其中一个为空就不再检查了
    if (node->left == nullptr || node->right == nullptr) {
        if (node->object->getBounds().IntersectP(ray, invDir, dirIsNeg)) {
            return node->object->getIntersection(ray);
        }
        return Intersection();
    }
    auto hit1 = getIntersection(node->left,ray);
    auto hit2 = getIntersection(node->right,ray);

    if (hit1.distance < hit2.distance) {
        return hit1;
    }
    return hit2;
}